#include "database.h"
#include <iostream>
#include <cstring>

Database::Database(const std::string& host, const std::string& user,
                   const std::string& password, const std::string& database, int port)
    : host_(host), user_(user), password_(password), database_(database), port_(port), conn_(nullptr) {
}

Database::~Database() {
    disconnect();
}

bool Database::connect() {
    if (conn_) {
        disconnect();
    }

    conn_ = mysql_init(nullptr);
    if (!conn_) {
        std::cerr << "mysql_init() failed" << std::endl;
        return false;
    }

    // 启用自动重连
    bool reconnect = 1;
    mysql_options(conn_, MYSQL_OPT_RECONNECT, &reconnect);

    if (!mysql_real_connect(conn_, host_.c_str(), user_.c_str(),
                           password_.c_str(), database_.c_str(), port_, nullptr, 0)) {
        std::cerr << "Database connection failed: " << mysql_error(conn_) << std::endl;
        mysql_close(conn_);
        conn_ = nullptr;
        return false;
    }

    // 设置字符集为UTF-8
    if (mysql_set_character_set(conn_, "utf8mb4") != 0) {
        std::cerr << "Setting character set failed: " << mysql_error(conn_) << std::endl;
        return false;
    }

    std::cout << "Database connected successfully" << std::endl;
    return true;
}

void Database::disconnect() {
    if (conn_) {
        mysql_close(conn_);
        conn_ = nullptr;
    }
}

bool Database::isConnected() const {
    return conn_ != nullptr;
}

std::unique_ptr<MYSQL_RES, decltype(&mysql_free_result)> Database::query(const std::string& sql) {
    std::lock_guard<std::mutex> lock(db_mutex_);

    if (!conn_) {
        std::cerr << "No database connection" << std::endl;
        return std::unique_ptr<MYSQL_RES, decltype(&mysql_free_result)>(nullptr, mysql_free_result);
    }

    // 清除之前的任何未读取的结果
    while (mysql_more_results(conn_)) {
        MYSQL_RES* dummy = mysql_store_result(conn_);
        if (dummy) mysql_free_result(dummy);
        mysql_next_result(conn_);
    }

    if (mysql_query(conn_, sql.c_str()) != 0) {
        std::cerr << "Query failed: " << mysql_error(conn_) << " (errno: " << mysql_errno(conn_) << ")" << std::endl;

        // 如果是连接错误，尝试重新连接
        if (mysql_errno(conn_) == CR_SERVER_LOST || mysql_errno(conn_) == CR_SERVER_GONE_ERROR ||
            mysql_errno(conn_) == CR_SERVER_LOST_EXTENDED || mysql_errno(conn_) == CR_CONN_HOST_ERROR) {
            std::cout << "Connection lost, attempting to reconnect..." << std::endl;
            disconnect();
            if (connect()) {
                // 重新连接成功，再次尝试查询
                if (mysql_query(conn_, sql.c_str()) == 0) {
                    MYSQL_RES* result = mysql_store_result(conn_);
                    return std::unique_ptr<MYSQL_RES, decltype(&mysql_free_result)>(result, mysql_free_result);
                }
            }
        }
        return std::unique_ptr<MYSQL_RES, decltype(&mysql_free_result)>(nullptr, mysql_free_result);
    }

    MYSQL_RES* result = mysql_store_result(conn_);
    if (!result && mysql_field_count(conn_) != 0) {
        std::cerr << "Failed to store result: " << mysql_error(conn_) << std::endl;
        return std::unique_ptr<MYSQL_RES, decltype(&mysql_free_result)>(nullptr, mysql_free_result);
    }

    return std::unique_ptr<MYSQL_RES, decltype(&mysql_free_result)>(result, mysql_free_result);
}

bool Database::execute(const std::string& sql) {
    std::lock_guard<std::mutex> lock(db_mutex_);

    if (!conn_) {
        std::cerr << "No database connection" << std::endl;
        return false;
    }

    // 清除之前的任何未读取的结果
    while (mysql_more_results(conn_)) {
        MYSQL_RES* dummy = mysql_store_result(conn_);
        if (dummy) mysql_free_result(dummy);
        mysql_next_result(conn_);
    }

    if (mysql_query(conn_, sql.c_str()) != 0) {
        std::cerr << "Execute failed: " << mysql_error(conn_) << " (errno: " << mysql_errno(conn_) << ")" << std::endl;

        // 如果是连接错误，尝试重新连接
        if (mysql_errno(conn_) == CR_SERVER_LOST || mysql_errno(conn_) == CR_SERVER_GONE_ERROR ||
            mysql_errno(conn_) == CR_SERVER_LOST_EXTENDED || mysql_errno(conn_) == CR_CONN_HOST_ERROR) {
            std::cout << "Connection lost, attempting to reconnect..." << std::endl;
            disconnect();
            if (connect()) {
                // 重新连接成功，再次尝试执行
                if (mysql_query(conn_, sql.c_str()) == 0) {
                    return true;
                }
            }
        }
        return false;
    }

    return true;
}

std::string Database::escapeString(const std::string& str) {
    std::lock_guard<std::mutex> lock(db_mutex_);

    if (!conn_) {
        return str;
    }

    char* escaped = new char[str.length() * 2 + 1];
    mysql_real_escape_string(conn_, escaped, str.c_str(), str.length());
    std::string result(escaped);
    delete[] escaped;
    return result;
}
